import warnings

import matplotlib
rcParams = matplotlib.rcParams
import matplotlib.artist as martist
import matplotlib.transforms as mtransforms
import matplotlib.collections as mcoll
import matplotlib.legend as mlegend

from matplotlib.axes import subplot_class_factory
try:
    from axislines import Axes
except:
    from matplotlib.axes import Axes


import matplotlib.axes as maxes

import matplotlib.cbook as cbook
is_string_like = cbook.is_string_like


class ParasiteAxes(Axes):

    def __init__(self, parent_axes, **kargs):

        self._parent_axes = parent_axes
        kargs.update(dict(frameon=False))
        super(ParasiteAxes, self).__init__(parent_axes.figure,
                                           parent_axes._position, **kargs)


    #def apply_aspect(self, position=None):
    #    pos = self._parent_axes.get_position(original=False)
    #    self.set_position(pos, "active")


    def cla(self):
        super(ParasiteAxes, self).cla()

        martist.setp(self.get_children(), visible=False)
        self._get_lines = self._parent_axes._get_lines


    def get_images_artists(self):
        artists = set([a for a in self.get_children() if a.get_visible()])
        images = set([a for a in self.images if a.get_visible()])

        return list(images), list(artists - images)



class ParasiteAxesAuxTrans(ParasiteAxes):

    def __init__(self, parent_axes, aux_transform, viewlim_mode=None):

        self.transAux = aux_transform

        self._viewlim_mode = viewlim_mode

        super(ParasiteAxesAuxTrans, self).__init__(parent_axes)

    def _set_lim_and_transforms(self):

        self.transAxes = self._parent_axes.transAxes

        self.transData = \
            self.transAux + \
            self._parent_axes.transData

        self._xaxis_transform = mtransforms.blended_transform_factory(
                self.transData, self.transAxes)
        self._yaxis_transform = mtransforms.blended_transform_factory(
                self.transAxes, self.transData)

    def update_viewlim(self):
        viewlim = self._parent_axes.viewLim.frozen()
        if self._viewlim_mode is None:
            pass
        elif self._viewlim_mode == "equal":
            self.axes.viewLim.set(viewlim)
        elif self._viewlim_mode == "transform":
            self.axes.viewLim.set(viewlim.transformed(self.transAux.inverted()))
        else:
            raise ValueError("Unknown mode : %s" % (self._viewlim_mode,))


    def apply_aspect(self, position=None):
        self.update_viewlim()
        super(ParasiteAxesAuxTrans, self).apply_aspect()


def _get_handles(ax):
    handles = ax.lines[:]
    handles.extend(ax.patches)
    handles.extend([c for c in ax.collections
                    if isinstance(c, mcoll.LineCollection)])
    handles.extend([c for c in ax.collections
                    if isinstance(c, mcoll.RegularPolyCollection)])
    return handles


class HostAxes(Axes):

    def __init__(self, *kl, **kwargs):

        self.parasites = []
        super(HostAxes, self).__init__(*kl, **kwargs)



    def legend(self, *args, **kwargs):

        if len(args)==0:
            all_handles = _get_handles(self)
            for ax in self.parasites:
                all_handles.extend(_get_handles(ax))
            handles = []
            labels = []
            for handle in all_handles:
                label = handle.get_label()
                if (label is not None and
                    label != '' and not label.startswith('_')):
                    handles.append(handle)
                    labels.append(label)
            if len(handles) == 0:
                warnings.warn("No labeled objects found. "
                              "Use label='...' kwarg on individual plots.")
                return None

        elif len(args)==1:
            # LABELS
            labels = args[0]
            handles = [h for h, label in zip(all_handles, labels)]

        elif len(args)==2:
            if is_string_like(args[1]) or isinstance(args[1], int):
                # LABELS, LOC
                labels, loc = args
                handles = [h for h, label in zip(all_handles, labels)]
                kwargs['loc'] = loc
            else:
                # LINES, LABELS
                handles, labels = args

        elif len(args)==3:
            # LINES, LABELS, LOC
            handles, labels, loc = args
            kwargs['loc'] = loc
        else:
            raise TypeError('Invalid arguments to legend')


        handles = cbook.flatten(handles)
        self.legend_ = mlegend.Legend(self, handles, labels, **kwargs)
        return self.legend_


    def draw(self, renderer):

        orig_artists = list(self.artists)
        orig_images = list(self.images)

        if hasattr(self, "get_axes_locator"):
            locator = self.get_axes_locator()
            if locator:
                pos = locator(self, renderer)
                self.set_position(pos, which="active")
                self.apply_aspect(pos)
            else:
                self.apply_aspect()
        else:
            self.apply_aspect()

        rect = self.get_position()

        for ax in self.parasites:
            ax.apply_aspect(rect)
            images, artists = ax.get_images_artists()
            self.images.extend(images)
            self.artists.extend(artists)

        super(HostAxes, self).draw(renderer)
        self.artists = orig_artists
        self.images = orig_images

    def cla(self):

        for ax in self.parasites:
            ax.cla()

        super(HostAxes, self).cla()


    def twinx(self):
        """
        call signature::

          ax = twinx()

        create a twin of Axes for generating a plot with a sharex
        x-axis but independent y axis.  The y-axis of self will have
        ticks on left and the returned axes will have ticks on the
        right
        """

        ax2 = ParasiteAxes(self, sharex=self, frameon=False)
        self.parasites.append(ax2)
        ax2.xaxis.set_visible(True)
        ax2.yaxis.set_visible(True)
        ax2.yaxis.tick_right()
        ax2.yaxis.set_label_position('right')
        self.yaxis.tick_left()
        return ax2

    def twiny(self):
        """
        call signature::

          ax = twiny()

        create a twin of Axes for generating a plot with a shared
        y-axis but independent x axis.  The x-axis of self will have
        ticks on bottom and the returned axes will have ticks on the
        top
        """

        ax2 = ParasiteAxes(self, sharey=self, frameon=False)
        self.parasites.append(ax2)
        ax2.xaxis.set_visible(True)
        ax2.yaxis.set_visible(True)
        ax2.xaxis.tick_top()
        ax2.xaxis.set_label_position('top')
        self.xaxis.tick_bottom()
        return ax2

SubplotHost = subplot_class_factory(HostAxes)


if __name__ == "__main__":
    from parasite_axes import HostAxes, ParasiteAxes
    from axislines import Axes
    import matplotlib.pyplot as plt
    fig = plt.figure(1)
    fig.clf()

    host = HostAxes(fig, [0.15, 0.1, 0.65, 0.8])
    par1 = ParasiteAxes(host, sharex=host)
    par2 = ParasiteAxes(host, sharex=host)
    host.parasites.append(par1)
    host.parasites.append(par2)

    host.set_ylabel("Density")
    host.set_xlabel("Distance")

    host.axislines["right"].set_visible(False)
    par1.axislines["right"].set_visible(True)
    par1.set_ylabel("Temperature")

    par1.axislines["right"].major_ticklabels.set_visible(True)
    par1.axislines["right"].label.set_visible(True)

    par2.set_ylabel("Velocity")
    offset_transform = mtransforms.Affine2D().translate(60, 0)
    axisline = par2._grid_helper.get_axisline
    par2.axislines["right2"] = axisline(loc="right",
                                        offset_transform=offset_transform)


    fig.add_axes(host)

    host.set_xlim(0, 2)
    host.set_ylim(0, 2)

    host.set_xlabel("Distance")
    host.set_ylabel("Density")
    par1.set_ylabel("Temperature")

    p1, = host.plot([0, 1, 2], [0, 1, 2], label="Density")
    p2, = par1.plot([0, 1, 2], [0, 3, 2], label="Temperature")
    p3, = par2.plot([0, 1, 2], [50, 30, 15], label="Velocity")

    par1.set_ylim(0, 4)
    par2.set_ylim(1, 65)

    host.legend()

    host.axislines["left"].label.set_color(p1.get_color())
    par1.axislines["right"].label.set_color(p2.get_color())
    par2.axislines["right2"].label.set_color(p3.get_color())

    plt.draw()

